#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
实时脑电状态分析演示程序

本文件演示如何将 eeg_data_logger 和 eeg_state_analyzer 集成，
实现实时脑电数据采集和状态分析。

功能特点：
1. 实时从NeuroSky设备获取脑电数据
2. 实时分析用户状态（清醒、闭眼放松、专注、困倦等）
3. 在控制台显示实时状态变化
4. 记录分析结果到日志文件

使用方法：
1. 确保NeuroSky设备已连接
2. 运行程序：python real_time_state_demo.py
3. 按Ctrl+C停止程序

原理：
- 使用NeuroPy库获取脑电数据
- 通过频段功率分析和注意力/冥想指数判断状态
- 采用滑动窗口平滑处理减少噪声影响
"""

import sys
import os
import time
import threading
from datetime import datetime

# 添加父目录到路径，以便导入模块
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
sys.path.append(os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'eeg'))

try:
    from eeg.eeg_data_logger import EEGDataLogger
    from eeg.eeg_state_analyzer import EEGStateAnalyzer
except ImportError as e:
    print(f"导入模块失败: {e}")
    print("请确保 eeg_data_logger.py 和 eeg_state_analyzer.py 文件存在")
    sys.exit(1)

class RealTimeStateDemo:
    """
    实时状态分析演示类
    
    集成数据采集和状态分析功能，提供实时状态监控
    """
    
    def __init__(self, log_file="real_time_analysis.csv"):
        """
        初始化演示程序
        
        Args:
            log_file: 日志文件名
        """
        self.log_file = log_file
        self.running = False
        self.current_state = "未知"
        self.confidence = 0.0
        self.last_update = None
        
        # 初始化数据记录器
        self.data_logger = EEGDataLogger(log_file=log_file)
        
        # 初始化状态分析器
        self.state_analyzer = EEGStateAnalyzer()
        
        # 数据缓存
        self.latest_data = {
            'attention': 0,
            'meditation': 0,
            'delta': 0,
            'theta': 0,
            'alpha': 0,
            'beta': 0,
            'gamma': 0,
            'signal_quality': 200
        }
        
        # 初始化neuropy连接
        self.data_logger.start_logging()
        
        # 设置回调函数
        self._setup_callbacks()
        
    def _setup_callbacks(self):
        """
        设置数据回调函数
        """
        # 重写数据记录器的回调函数，添加状态分析
        original_attention_callback = self.data_logger.attention_callback
        original_meditation_callback = self.data_logger.meditation_callback
        
        def enhanced_attention_callback(attention_value):
            self.latest_data['attention'] = attention_value
            original_attention_callback(attention_value)
            self._analyze_state()
            
        def enhanced_meditation_callback(meditation_value):
            self.latest_data['meditation'] = meditation_value
            original_meditation_callback(meditation_value)
            self._analyze_state()
            
        # 频段功率回调
        def enhanced_delta_callback(delta_value):
            self.latest_data['delta'] = delta_value
            self.data_logger.delta_callback(delta_value)
            
        def enhanced_theta_callback(theta_value):
            self.latest_data['theta'] = theta_value
            self.data_logger.theta_callback(theta_value)
            
        def enhanced_alpha_callback(alpha_value):
            self.latest_data['alpha'] = alpha_value
            self.data_logger.alpha_callback(alpha_value)
            
        def enhanced_beta_callback(beta_value):
            self.latest_data['beta'] = beta_value
            self.data_logger.beta_callback(beta_value)
            
        def enhanced_gamma_callback(gamma_value):
            self.latest_data['gamma'] = gamma_value
            self.data_logger.gamma_callback(gamma_value)
            
        def enhanced_signal_quality_callback(signal_quality):
            self.latest_data['signal_quality'] = signal_quality
            self.data_logger.signal_quality_callback(signal_quality)
            
        # 替换回调函数
        self.data_logger.neuropy.setCallBack("attention", enhanced_attention_callback)
        self.data_logger.neuropy.setCallBack("meditation", enhanced_meditation_callback)
        self.data_logger.neuropy.setCallBack("lowAlpha", enhanced_alpha_callback)
        self.data_logger.neuropy.setCallBack("highAlpha", enhanced_alpha_callback)
        self.data_logger.neuropy.setCallBack("lowBeta", enhanced_beta_callback)
        self.data_logger.neuropy.setCallBack("highBeta", enhanced_beta_callback)
        self.data_logger.neuropy.setCallBack("lowGamma", enhanced_gamma_callback)
        self.data_logger.neuropy.setCallBack("highGamma", enhanced_gamma_callback)
        self.data_logger.neuropy.setCallBack("delta", enhanced_delta_callback)
        self.data_logger.neuropy.setCallBack("theta", enhanced_theta_callback)
        self.data_logger.neuropy.setCallBack("poorSignalLevel", enhanced_signal_quality_callback)
        
    def _analyze_state(self):
        """
        分析当前状态
        """
        # 只有在信号质量良好时才进行分析
        if self.latest_data['signal_quality'] > 100:
            return
            
        # 进行状态分析
        state, confidence = self.state_analyzer.analyze_state(
            attention=self.latest_data['attention'],
            meditation=self.latest_data['meditation'],
            delta=self.latest_data['delta'],
            theta=self.latest_data['theta'],
            alpha=self.latest_data['alpha'],
            beta=self.latest_data['beta'],
            gamma=self.latest_data['gamma'],
            signal_quality=self.latest_data['signal_quality']
        )
        
        # 更新状态
        if state != self.current_state or abs(confidence - self.confidence) > 0.1:
            self.current_state = state
            self.confidence = confidence
            self.last_update = datetime.now()
            self._display_status()
            
    def _display_status(self):
        """
        显示当前状态
        """
        timestamp = self.last_update.strftime("%H:%M:%S")
        signal_status = "良好" if self.latest_data['signal_quality'] < 50 else "一般" if self.latest_data['signal_quality'] < 100 else "差"
        
        print(f"\r[{timestamp}] 状态: {self.current_state:8s} | "
              f"置信度: {self.confidence:.2f} | "
              f"注意力: {self.latest_data['attention']:2d} | "
              f"冥想: {self.latest_data['meditation']:2d} | "
              f"信号: {signal_status}", end="", flush=True)
              
    def start(self):
        """
        开始实时分析
        """
        print("=" * 60)
        print("实时脑电状态分析演示")
        print("=" * 60)
        print("正在连接NeuroSky设备...")
        
        try:
            # 启动数据记录
            self.data_logger.start_recording()
            self.running = True
            
            print("设备连接成功！开始实时分析...")
            print("按 Ctrl+C 停止程序")
            print("-" * 60)
            
            # 主循环
            while self.running:
                time.sleep(0.1)
                
        except KeyboardInterrupt:
            print("\n\n用户中断程序")
        except Exception as e:
            print(f"\n\n程序出错: {e}")
        finally:
            self.stop()
            
    def stop(self):
        """
        停止分析
        """
        print("\n正在停止程序...")
        self.running = False
        
        try:
            self.data_logger.stop_recording()
            print("数据记录已停止")
        except:
            pass
            
        print(f"分析结果已保存到: {self.log_file}")
        print("程序已退出")

def main():
    """
    主函数
    """
    # 创建演示实例
    demo = RealTimeStateDemo("real_time_analysis.csv")
    
    # 开始演示
    demo.start()

if __name__ == "__main__":
    main()